import torch
import torch.nn     as nn
from   transformers import CLIPTextModel

#! 1. 词嵌入层部分
class Embedding(nn.Module):
    def __init__(self,):
        super(Embedding, self).__init__()
        self.embedding     = nn.Embedding(49408, 768)
        self.pos_embedding = nn.Embedding(77, 768)
        """
        @brief register_buffer是一个非常有用的方法，它用于在模块中注册一个缓冲区。这些缓冲区是模块的一部分，但并不是
            模型参数的一部分，也就是说，它们在训练过程中通常不会被优化器更新。然而，它们对模型的运行是必要的，因为它们可
            能携带有对于模型运行逻辑必要的数据。使用register_buffer可以确保这些数据能够随着模型一起被保存和加载，而无
            需进行梯度计算。
        """
        self.register_buffer("pos_ids", torch.arange(77).unsqueeze(dim=0)) # [1, 77]
        
    def forward(self, input_ids:torch.Tensor):
        embed     = self.embedding(input_ids)         # [bs, 77] ==> [bs, 77, 768]
        pos_embed = self.pos_embedding(self.pos_ids)  # [1, 77]  ==> [1,  77, 768]
        return embed + pos_embed                      # [bs, 77, 768]


#! 2. Transformer模型的多头注意力机制部分    
class MHA(nn.Module):
    def __init__(self):
        super(MHA, self).__init__()
        self.Q = nn.Linear(768, 768)
        self.K = nn.Linear(768, 768)
        self.V = nn.Linear(768, 768)
        self.O = nn.Linear(768, 768)                  # out
        
    # 覆盖mask
    def get_mask():
        """
        @brief 5x5示例
            tensor([[0., -inf, -inf, -inf, -inf],
                [0., 0., -inf, -inf, -inf],
                [0., 0., 0., -inf, -inf],
                [0., 0., 0., 0., -inf],
                [0., 0., 0., 0., 0.]])
        """
        mask = torch.empty(77, 77)
        mask.fill_(-float('inf'))                     # 所有值设置为负无穷
        mask.triu_(1)                                 # 对角线和以下的位置为0
        return mask
    
    def forward(self, x:torch.Tensor):                # [bs, 77, 768]   77是max_seq_len
        b = x.shape[0]            
        q = self.Q(x) * 0.125                         # [bs, 77, 768]     
        k = self.K(x)
        v = self.V(x)
        
        #! 拆分注意力头
        #  [b, 77, 768] -> [b, 77, 12, 64] -> [b, 12, 77, 64] -> [b*12, 77, 64]
        q = q.reshape(-1, 77, 12, 64).transpose(1, 2).reshape(-1, 77, 64)
        k = k.reshape(-1, 77, 12, 64).transpose(1, 2).reshape(-1, 77, 64)
        v = v.reshape(-1, 77, 12, 64).transpose(1, 2).reshape(-1, 77, 64)
        
        # 计算qk乘积: [b*12, 77, 64] * [b*12, 64, 77] -> [b*12, 77, 77]
        attn = torch.bmm(q, k.transpose(1, 2))        
        attn = attn + self.get_mask()
        attn = attn.softmax(dim=-1)                   # 计算softmax,被mask的部分值为0
        attn = torch.bmm(attn, v)                     # 计算和v的乘积[b*12, 77, 77] * [b*12, 77, 64] -> [b*12, 77, 64]
    
        # [b*12, 77, 64] -> [b, 12, 77, 64] -> [b, 77, 12, 64] -> [b, 77, 768]
        attn = attn.reshape(-1, 12, 77, 64).transpose(1, 2).reshape(-1, 77, 768)
        return self.out(attn)                         # 线性输出,维度不变 [b, 77, 768]
    
    
#! 3. Tranformer模型的EncoderLayer
class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.seq1 = nn.Sequential(
                nn.LayerNorm(768),
                MHA(),
            )
        
        self.seq2 = nn.Sequential(
                nn.LayerNorm(768),
                nn.Linear(768, 3072),
            )
                
        self.seq3 = nn.Linear(3072, 768)
        
    def forward(self, x:torch.Tensor):                 # x -> [2, 77, 768]
        x        = x + self.seq1(x)
        shortcut = x
        x        = self.seq2(x)
        x        = x * (x * 1.702).sigmoid()           #! TODO: quick gelu的激活函数
        return shortcut + self.seq3(x)
    
    
def build_encoder():
    encoder = nn.Sequential(
            Embedding(),
            EncoderLayer(), EncoderLayer(), EncoderLayer(), EncoderLayer(),
            EncoderLayer(), EncoderLayer(), EncoderLayer(), EncoderLayer(),
            EncoderLayer(), EncoderLayer(), EncoderLayer(), EncoderLayer(),
            nn.LayerNorm(768)
        )
    params = CLIPTextModel.from_pretrained('lansinuote/diffsion_from_scratch.params', subfolder='text_encoder')      # 加载预训练模型的参数
    encoder[0].embed.load_state_dict(params.text_model.embeddings.token_embedding.state_dict())                      # 词编码
    encoder[0].pos_embed.load_state_dict(params.text_model.embeddings.position_embedding.state_dict())               # 位置编码

    #12层编码层
    for i in range(12):
        encoder[i + 1].seq1[0].load_state_dict(params.text_model.encoder.layers[i].layer_norm1.state_dict())         # 第一层norm
        encoder[i + 1].seq1[1].q.load_state_dict(params.text_model.encoder.layers[i].self_attn.q_proj.state_dict())  # 注意力q矩阵
        encoder[i + 1].seq1[1].k.load_state_dict(params.text_model.encoder.layers[i].self_attn.k_proj.state_dict())  # 注意力k矩阵
        encoder[i + 1].seq1[1].v.load_state_dict(params.text_model.encoder.layers[i].self_attn.v_proj.state_dict())  # 注意力v矩阵

        encoder[i + 1].seq1[1].out.load_state_dict(params.text_model.encoder.layers[i].self_attn.out_proj.state_dict())  # 注意力out

        encoder[i + 1].seq2[0].load_state_dict(params.text_model.encoder.layers[i].layer_norm2.state_dict())         # 第二层norm
        encoder[i + 1].seq2[1].load_state_dict(params.text_model.encoder.layers[i].mlp.fc1.state_dict())             # mlp第一层fc
        encoder[i + 1].seq3.load_state_dict(params.text_model.encoder.layers[i].mlp.fc2.state_dict())                # mlp第二层fc
    encoder[13].load_state_dict(params.text_model.final_layer_norm.state_dict())                                     # 输出norm
    
    return encoder


